#!/usr/bin/env python3

# Copyright (c) 2005-2014 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import enum
import anytree
from ._util import Util
from ._util import PhysicalDiskMounts
from ._exception import RunningEnvironmentError


class Flags(enum.Flag):
    BIOS = enum.auto()
    BOOTLOADER = enum.auto()
    KERNEL = enum.auto()
    INITRAMFS = enum.auto()


class KernelType(enum.Enum):
    LINUX = "linux"


class Arch(enum.Enum):
    X86_64 = "x86_64"


class BootMode(enum.Enum):
    EFI = "efi"
    BIOS = "bios"


class SystemInit:

    class Type(enum.Enum):
        SYSVINIT = "sysv-init"
        OPENRC = "openrc"
        SYSTEMD = "systemd"
        CUSTOM = "custom"

    def __init__(self, system_init_type, cmd=None):
        self.system_init_type = system_init_type
        if self.system_init_type == self.Type.SYSVINIT:
            # FIXME
            assert False
        elif self.system_init_type == self.Type.OPENRC:
            assert cmd is None
            self.cmd = "/sbin/openrc-init"
        elif self.system_init_type == self.Type.SYSTEMD:
            assert cmd is None
            self.cmd = "/usr/lib/systemd/systemd"
        elif self.system_init_type == self.Type.CUSTOM:
            assert cmd is not None
            self.cmd = cmd
        else:
            assert False

    def __eq__(self, other):
        if type(self) is not type(other):
            return False

        if self.system_init_type != other.name:
            return False

        if self.system_init_type == self.Type.SYSVINIT:
            assert self.cmd == other.cmd
        elif self.system_init_type == self.Type.OPENRC:
            assert self.cmd == other.cmd
        elif self.system_init_type == self.Type.SYSTEMD:
            assert self.cmd == other.cmd
        elif self.system_init_type == self.Type.CUSTOM:
            if self.cmd != other.cmd:
                return False
        else:
            assert False

        return True

    def __hash__(self):
        return hash((self.system_init_type, self.cmd))


class RescueOsSpec:

    def __init__(self, bbki):
        self.root_dir = bbki._fsLayout.get_boot_rescue_os_dir()
        self.kernel_filepath = os.path.join(self.root_dir, "vmlinuz")
        self.kernel_cmdline_filepath = os.path.join(self.root_dir, "vmlinuz.cmdline")
        self.initrd_filepath = os.path.join(self.root_dir, "initrd.img")

    def __repr__(self):
        return "<%s %r>" % (self.__class__.__name__, self.__dict__)


class HostMountPoint:

    FS_TYPE_VFAT = "vfat"
    FS_TYPE_EXT4 = "ext4"
    FS_TYPE_BTRFS = "btrfs"
    FS_TYPE_BCACHEFS = "bcachefs"

    def __init__(self, mount_point, dev_path_or_uuid, fs_type=None, mnt_opts=None, underlay_disk=None):
        self.device = None              # use same variable name as the namedtuple elements in psutil.disk_partitions()
        self.mountpoint = None          # same as above
        self.fstype = None              # same as above
        self.opts = None                # same as above
        self.dev_uuid = None
        self.underlay_disk = None

        # self.mountpoint
        assert os.path.isabs(mount_point)
        self.mountpoint = mount_point

        # self.device and self.dev_uuid, may contain multiple values seperated by ":"
        if ":" not in dev_path_or_uuid:
            if dev_path_or_uuid.startswith("/dev/"):
                self.device = dev_path_or_uuid
                self.dev_uuid = "UUID=%s" % (Util.getBlkDevUuid(dev_path_or_uuid))
            elif re.fullmatch(r'UUID=\S+', dev_path_or_uuid):
                self.device = None
                self.dev_uuid = dev_path_or_uuid
            else:
                assert False
        else:
            tlist = dev_path_or_uuid.split(":")
            if all([item.startswith("/dev/") for item in tlist]):
                self.device = dev_path_or_uuid
                self.dev_uuid = ":".join(["UUID_SUB=%s" % (Util.getBlkDevSubUuid(item)) for item in tlist])
            elif all([re.fullmatch(r'UUID_SUB=\S+', item) for item in tlist]):
                self.device = None
                self.dev_uuid = dev_path_or_uuid
            else:
                assert False

        # self.fstype
        if fs_type is not None:
            assert fs_type in [self.FS_TYPE_VFAT, self.FS_TYPE_EXT4, self.FS_TYPE_BTRFS, self.FS_TYPE_BCACHEFS]
            self.fstype = fs_type
        else:
            assert self.device is not None
            for item in self.device.split(":"):
                self.fstype = Util.getBlkDevFsType(item)
                if self.fstype not in [self.FS_TYPE_VFAT, self.FS_TYPE_EXT4, self.FS_TYPE_BTRFS, self.FS_TYPE_BCACHEFS]:
                    raise RunningEnvironmentError("device \"%s\" has unsupported filesystem %s" % (item, self.fstype))

        # self.opts
        if mnt_opts is not None:
            assert fs_type is not None
            assert isinstance(mnt_opts, str)
            self.opts = mnt_opts
        else:
            assert fs_type is None
            assert self.device is not None
            self.opts = PhysicalDiskMounts.find_entry_by_mount_point(self.mountpoint).opts

        # self.underlay_disk
        if underlay_disk is not None:
            assert fs_type is not None and mnt_opts is not None
            assert all([isinstance(x, HostDisk) for x in anytree.PostOrderIter(underlay_disk)])
            self.underlay_disk = underlay_disk
        else:
            assert self.device is not None
            self.underlay_disk = HostDisk.getUnderlayDisk(self.device, mount_point=self.mountpoint)

    def __eq__(self, other):
        if type(self) is not type(other):
            return False
        if self.device != other.device:
            return False
        if self.mountpoint != other.mountpoint:
            return False
        if self.fstype != other.fstype:
            return False
        if self.opts != other.opts:
            return False
        if self.dev_uuid != other.dev_uuid:
            return False
        return True

    def __repr__(self):
        return "<%s %r>" % (self.__class__.__name__, self.__dict__)


class HostDisk(anytree.node.nodemixin.NodeMixin):

    def __init__(self, uuid, parent):
        assert re.fullmatch(r'(UUID=|UUID_SUB=)\S+', uuid)
        super().__init__()
        self.parent = parent
        self.uuid = uuid

    def __eq__(self, other):
        return type(self) is type(other) and self.uuid == other.uuid

    def __hash__(self):
        return hash(self.uuid)

    @staticmethod
    def getUnderlayDisk(devPath, parent=None, mount_point=None):
        if parent is None:
            assert mount_point is not None
        else:
            assert mount_point is None

        klassList = [
            HostDiskBtrfsRaid,
            HostDiskBcachefsRaid,
            HostDiskLvmLv,
            HostDiskBcache,
            HostDiskScsiHdd,
            HostDiskXenHdd,
            HostDiskVirtioHdd,
            HostDiskNvmeHdd,
        ]
        for klass in klassList:
            ret = klass.getUnderlayDisk(devPath, parent, mount_point)
            if ret is not None:
                return ret

        # unknown
        raise RunningEnvironmentError("unknown device \"%s\"" % (devPath))

    @staticmethod
    def _getSubUuidOrUuidWithPrefix(devPath):
        ret = Util.getBlkDevSubUuid(devPath)
        if ret is not None:
            return "UUID_SUB=" + ret
        ret = Util.getBlkDevUuid(devPath)
        if ret is not None:
            return "UUID=" + ret
        assert False


class HostDiskWholeDiskOrPartition(HostDisk):

    WHOLE_DISK = 1
    MBR_PARTITION = 2
    GPT_PARTITION = 3

    def __init__(self, uuid, partition_type, parent):
        super().__init__(uuid, parent)
        assert partition_type in [self.WHOLE_DISK, self.MBR_PARTITION, self.GPT_PARTITION]
        self.partition_type = partition_type

    @classmethod
    def _getPartitionType(cls, diskDevPath, partitionId):
        if partitionId is None:
            return cls.WHOLE_DISK
        else:
            m2 = re.search(r'(^| )PTTYPE="(\S+)"', Util.cmdCall("blkid", diskDevPath), re.M)
            if m2.group(2) == "dos":
                return cls.MBR_PARTITION
            elif m2.group(2) == "gpt":
                return cls.GPT_PARTITION
            else:
                raise RunningEnvironmentError("unknown partition type \"%s\" for block device \"%s\"" % (m2.group(2), diskDevPath))


class HostDiskBtrfsRaid(HostDisk):

    def __init__(self, uuid, parent):
        # uuid: FS-UUID of the whole btrfs filesystem
        super().__init__(uuid, parent)

    @classmethod
    def getUnderlayDisk(cls, devPath, parent, mountPoint):
        if os.path.exists(devPath):
            if parent is None and Util.getBlkDevFsType(devPath) == "btrfs":
                bdi = cls("UUID=" + Util.btrfsGetUuid(mountPoint), parent)
                for slaveDevPath in Util.btrfsGetSlavePathList(mountPoint):
                    HostDisk.getUnderlayDisk(slaveDevPath, parent=bdi)
                return bdi
        return None


class HostDiskBcachefsRaid(HostDisk):

    def __init__(self, uuid, parent):
        # uuid: FS-UUID of the whole bcachefs filesystem
        super().__init__(uuid, parent)

    @classmethod
    def getUnderlayDisk(cls, devPath, parent, mountPoint):
        if mountPoint is not None:
            slaveDevPathList = devPath.split(":")
            if len(slaveDevPathList) > 1:
                if any([Util.getBlkDevFsType(x) == "bcachefs" for x in slaveDevPathList]):
                    assert all([Util.getBlkDevFsType(x) == "btrfs" for x in slaveDevPathList])
                    bdi = cls("UUID=" + Util.bcachefsGetUuid(slaveDevPathList), parent)
                    for slaveDevPath in slaveDevPathList:
                        HostDisk.getUnderlayDisk(slaveDevPath, parent=bdi)
                    return bdi
        return None


class HostDiskLvmLv(HostDisk):

    def __init__(self, uuid, vg_name, lv_name, parent):
        # uuid: FS-UUID of the filesystem in this LV, or SUB-UUID takes priority for multi-volume filesystem
        super().__init__(uuid, parent)
        self.vg_name = vg_name
        self.lv_name = lv_name

    @classmethod
    def getUnderlayDisk(cls, devPath, parent, mountPoint):
        lvmInfo = Util.getBlkDevLvmInfo(devPath)
        if lvmInfo is not None:
            assert mountPoint is not None
            bdi = cls(cls._getSubUuidOrUuidWithPrefix(devPath), lvmInfo[0], lvmInfo[1], parent)
            for slaveDevPath in Util.lvmGetSlaveDevPathList(lvmInfo[0]):
                HostDisk.getUnderlayDisk(slaveDevPath, parent=bdi)
            return bdi
        return None


class HostDiskBcache(HostDisk):

    def __init__(self, uuid, parent):
        # uuid: FS-UUID of the filesystem in this bcache device, or SUB-UUID takes priority for multi-volume filesystem
        super().__init__(uuid, parent)
        self.cache_dev_list = []
        self.backing_dev = None

    def add_cache_dev(self, disk):
        self.cache_dev_list.append(disk)

    def add_backing_dev(self, disk):
        assert self.backing_dev is None
        self.backing_dev = disk

    @classmethod
    def getUnderlayDisk(cls, devPath, parent, mountPoint):
        m = re.fullmatch("/dev/bcache[0-9]+", devPath)
        if m is not None:
            bdi = cls(cls._getSubUuidOrUuidWithPrefix(devPath), parent)
            slist = Util.bcacheGetSlaveDevPathList(devPath)
            for i in range(0, len(slist)):
                if i < len(slist) - 1:
                    bdi.add_cache_dev(HostDisk.getUnderlayDisk(slist[i], parent=bdi))
                else:
                    bdi.add_backing_dev(HostDisk.getUnderlayDisk(slist[i], parent=bdi))
            return bdi
        return None


class HostDiskScsiHdd(HostDiskWholeDiskOrPartition):

    def __init__(self, uuid, partition_type, host_controller_name, parent):
        # uuid: FS-UUID of the filesystem in this bcache device, or SUB-UUID takes priority for multi-volume filesystem
        super().__init__(uuid, partition_type, parent)
        self.host_controller_name = host_controller_name

    @classmethod
    def getUnderlayDisk(cls, devPath, parent, mountPoint):
        m = re.fullmatch("(/dev/sd[a-z])([0-9]+)?", devPath)
        if m is not None:
            ptype = cls._getPartitionType(m.group(1), m.group(2))
            return cls(cls._getSubUuidOrUuidWithPrefix(devPath), ptype, Util.scsiGetHostControllerName(m.group(1)), parent)
        return None


class HostDiskNvmeHdd(HostDiskWholeDiskOrPartition):

    def __init__(self, uuid, partition_type, parent):
        # uuid: FS-UUID of the filesystem in this bcache device, or SUB-UUID takes priority for multi-volume filesystem
        super().__init__(uuid, partition_type, parent)

    @classmethod
    def getUnderlayDisk(cls, devPath, parent, mountPoint):
        m = re.fullmatch("(/dev/nvme[0-9]+n[0-9]+)(p([0-9]+))?", devPath)
        if m is not None:
            ptype = cls._getPartitionType(m.group(1), m.group(3))
            return cls(cls._getSubUuidOrUuidWithPrefix(devPath), ptype, parent)
        return None


class HostDiskXenHdd(HostDiskWholeDiskOrPartition):

    def __init__(self, uuid, partition_type, parent):
        # uuid: FS-UUID of the filesystem in this bcache device, or SUB-UUID takes priority for multi-volume filesystem
        super().__init__(uuid, partition_type, parent)

    @classmethod
    def getUnderlayDisk(cls, devPath, parent, mountPoint):
        m = re.fullmatch("(/dev/xvd[a-z])([0-9]+)?", devPath)
        if m is not None:
            ptype = cls._getPartitionType(m.group(1), m.group(2))
            return cls(cls._getSubUuidOrUuidWithPrefix(devPath), ptype, parent)
        return None


class HostDiskVirtioHdd(HostDiskWholeDiskOrPartition):

    def __init__(self, uuid, parent):
        # uuid: FS-UUID of the filesystem in this bcache device, or SUB-UUID takes priority for multi-volume filesystem
        super().__init__(uuid, parent)

    @classmethod
    def getUnderlayDisk(cls, devPath, parent, mountPoint):
        m = re.fullmatch("(/dev/vd[a-z])([0-9]+)?", devPath)
        if m is not None:
            ptype = cls._getPartitionType(m.group(1), m.group(2))
            return cls(cls._getSubUuidOrUuidWithPrefix(devPath), ptype, parent)
        return None


class AuxOs:

    class Type(enum.Enum):
        MICROSOFT_WINDOWS = "microsoft-windows"

    def __init__(self, name, os_type, partition_path_or_uuid, chainloader_number):
        assert isinstance(os_type, self.Type)

        self.name = name
        self.os_type = os_type
        if partition_path_or_uuid.startswith("/dev/"):
            self.partition_path = partition_path_or_uuid
            self.partition_uuid = "UUID=" + Util.getBlkDevUuid(self.partition_path)
        else:
            self.partition_path = None
            self.partition_uuid = partition_path_or_uuid
        self.chainloader_number = chainloader_number

    def get_paths_in_boot_dir(self):
        if self.os_type == self.OS_TYPE_MICROSOFT_WINDOWS:
            return [
                "EFI/Microsoft",
                "McAfee",
            ]
        else:
            assert False

    def __eq__(self, other):
        if type(self) is not type(other):
            return False
        if self.name != other.name:
            return False
        if self.os_type != other.os_type:
            return False
        if self.partition_uuid != other.partition_uuid:
            return False
        if self.chainloader_number != other.chainloader_number:
            return False
        if self.partition_path is not None and other.partition_path is not None:
            assert self.partition_path == other.partition_path
        return True

    def __hash__(self):
        return hash((self.name, self.os_type, self.partition_uuid, self.chainloader_number))


class FsLayout:

    def get_boot_dir(self):
        return "/boot"

    def get_boot_efi_dir(self):
        return os.path.join(self.get_boot_dir(), "EFI")

    def get_boot_grub_dir(self):
        return os.path.join(self.get_boot_dir(), "grub")

    def get_boot_rescue_os_dir(self):
        return os.path.join(self.get_boot_dir(), "rescue")

    def get_boot_history_dir(self):
        return os.path.join(self.get_boot_dir(), "history")

    def get_kernel_modules_dir(self, kernel_verstr=None):
        if kernel_verstr is None:
            return "/lib/modules"
        else:
            return "/lib/modules/%s" % (kernel_verstr)

    def get_firmware_dir(self):
        return "/lib/firmware"

    def get_bbki_cache_dir(self):
        return "/var/cache/bbki"

    def get_bbki_cache_versions_dir(self):
        return os.path.join(self.get_bbki_cache_dir(), "versions")

    def get_bbki_cache_distfiles_dir(self):
        return os.path.join(self.get_bbki_cache_dir(), "distfiles")

    def get_bbki_cache_distfiles_ro_dir_list(self):
        return []

    def get_bbki_tmp_dir(self):
        return "/var/tmp/bbki"
